package main

import (
	"crypto/md5"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/jinzhu/gorm"
	_ "github.com/jinzhu/gorm/dialects/sqlite"
	"github.com/leaanthony/mewn"
)

var streamerMap = make(map[string]*streamer)
var streamerMapLock sync.RWMutex

var db *gorm.DB

var port int
var ffmpegPath string

var workDir string

var tokenMap map[string]*token = make(map[string]*token)
var tokenMapLock sync.RWMutex
var rootUser string = "admin"
var rootPwd string = "123"

// go build -ldflags="-w -s"
// go build -ldflags="-w -s -H windowsgui"
// ./mewn build -ldflags="-w -s"
// ./mewn build -ldflags="-w -s -H windowsgui"
// ./upx -9 ./jsmpeg-streamer
// ./upx.exe -9 ./jsmpeg-streamer.exe
func main() {
	root, err := filepath.Abs(filepath.Dir(os.Args[0]))
	if err != nil {
		log.Fatal(err)
	}
	workDir = root
	log.Println("work dir " + workDir)

	logFile, err := os.OpenFile(filepath.Join(workDir, "jsmpeg-streamer.log"), os.O_RDWR|os.O_CREATE|os.O_APPEND, 0766)
	if err != nil {
		log.Fatal(err)
	}
	defer logFile.Close()
	w := io.MultiWriter(logFile, os.Stdout)
	log.SetOutput(w)
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	flag.IntVar(&port, "p", 10019, "端口, 默认10019")
	flag.StringVar(&ffmpegPath, "ff", "", "ffmpeg路径, 默认同级目录或环境变量")
	flag.Parse()

	log.Println("using port " + strconv.Itoa(port))

	ffmpegCheck()

	db, err = gorm.Open("sqlite3", filepath.Join(root, "jsmpeg-streamer.db?cache=shared"))
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()
	db.DB().SetMaxOpenConns(1)

	db.AutoMigrate(&streamer{})
	var streamers []*streamer
	if err := db.Find(&streamers).Error; err != nil {
		log.Fatalln(err.Error())
	}
	for _, s := range streamers {
		s.Alive = false
		s.PlayerMap = make(map[string]*player)
		s.PlayerMapLock = &sync.RWMutex{}
		streamerMapLock.Lock()
		streamerMap[s.Key] = s
		streamerMapLock.Unlock()
		if !s.Lazy {
			go s.run()
		}
	}

	db.AutoMigrate(&user{})
	var u user
	db.Where(&user{
		Username: rootUser,
	}).First(&u)
	if u.Username == "" {
		u.Username = rootUser
		u.Password = rootPwd
		err = db.Create(&u).Error
		if err != nil {
			log.Fatalln(err)
		}
	} else {
		rootUser = u.Username
		rootPwd = u.Password
	}

	r := gin.Default()
	r.Use(func(context *gin.Context) {
		context.Writer.Header().Set("Access-Control-Allow-Origin", "*")
		context.Next()
	})

	// WebSocket realy
	r.GET("/relay", realy)

	r.POST("/api/streamer/add", streamerAdd)
	r.GET("/api/streamer/list", streamerList)
	r.POST("/api/streamer/delete", streamerDelete)
	r.POST("/api/user/login", userLogin)
	r.POST("/api/user/logout", userLogout)
	r.POST("/api/user/resetpwd", userResetPwd)

	// static resources
	indexHTML := mewn.Bytes("./www/index.html")
	favicon := mewn.Bytes("./www/favicon.png")
	basksJPG := mewn.Bytes("./www/static/img/basks.jpg")
	appCSS := mewn.Bytes("./www/static/css/app.css")
	//appCSSMap := mewn.Bytes("./www/static/css/app.css.map")
	elementIconsTtf := mewn.Bytes("./www/static/fonts/element-icons.ttf")
	elementIconsWoff := mewn.Bytes("./www/static/fonts/element-icons.woff")
	appJs := mewn.Bytes("./www/static/js/app.js")
	//appJsMap := mewn.Bytes("./www/static/js/app.js.map")
	manifestJs := mewn.Bytes("./www/static/js/manifest.js")
	//manifestJsMap := mewn.Bytes("./www/static/js/manifest.js.map")
	vendorJs := mewn.Bytes("./www/static/js/vendor.js")
	//vendorJsMap := mewn.Bytes("./static/www/js/vendor.js.map")

	r.GET("/", func(c *gin.Context) {
		c.Header("Content-Type", "text/html")
		c.Writer.Write(indexHTML)
	})
	r.GET("/index.html", func(c *gin.Context) {
		c.Header("Content-Type", "text/html")
		c.Writer.Write(indexHTML)
	})
	r.GET("/favicon.png", func(c *gin.Context) {
		c.Header("Content-Type", "image/x-icon")
		c.Writer.Write(favicon)
	})
	r.GET("/static/css/app.css", func(c *gin.Context) {
		c.Header("Content-Type", "text/css")
		c.Writer.Write(appCSS)
	})
	r.GET("/static/fonts/element-icons.ttf", func(c *gin.Context) {
		c.Writer.Write(elementIconsTtf)
	})
	r.GET("/static/fonts/element-icons.woff", func(c *gin.Context) {
		c.Writer.Write(elementIconsWoff)
	})
	r.GET("/static/js/app.js", func(c *gin.Context) {
		c.Header("Content-Type", "text/javascript")
		c.Writer.Write(appJs)
	})
	r.GET("/static/js/manifest.js", func(c *gin.Context) {
		c.Header("Content-Type", "text/javascript")
		c.Writer.Write(manifestJs)
	})
	r.GET("/static/js/vendor.js", func(c *gin.Context) {
		c.Header("Content-Type", "text/javascript")
		c.Writer.Write(vendorJs)
	})
	r.GET("/static/img/basks.jpg", func(c *gin.Context) {
		c.Header("Content-Type", "image/jpeg")
		c.Writer.Write(basksJPG)
	})

	go func() {
		for {
			for _, s := range streamerMap {
				if s != nil {
					if !s.Lazy {
						if !s.Alive {
							log.Println("!lazy streamer " + s.Key + " dropped, try start")
							go s.run()
						}
					} else {
						if len(s.PlayerMap) == 0 && s.Alive && time.Now().Unix() > s.LastDisconnectTime.Unix()+60 {
							s.stop()
							log.Println("streamer " + s.Key + " has no player, stop")
						}
					}
				}
			}
			for key, t := range tokenMap {
				if t != nil && time.Now().Unix() > t.Time.Unix()+60*60*24 {
					tokenMapLock.Lock()
					delete(tokenMap, key)
					tokenMapLock.Unlock()
				}
			}
			time.Sleep(time.Second * 10)
		}
	}()

	gin.SetMode(gin.ReleaseMode)
	r.Run(":" + strconv.Itoa(port))
}

func ffmpegCheck() {
	if ffmpegPath == "" {
		log.Println("未指定ffmepg路径, 查找同级目录")
		if runtime.GOOS == "windows" {
			ffmpegPath = workDir + string(os.PathSeparator) + "ffmpeg.exe"
		} else {
			ffmpegPath = workDir + string(os.PathSeparator) + "ffmpeg"
		}
		exists := pathExist(ffmpegPath)
		if exists {
			log.Println("使用同级目录: " + ffmpegPath)
		} else {
			ffmpegPath = "ffmpeg"
			log.Println("同级目录未找到")
			log.Println("使用环境变量: " + ffmpegPath)
		}
	}
	// try exec ffmpeg
	log.Println("测试执行 ffmpeg -version")
	cmd := exec.Command(ffmpegPath, "-version")
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Fatal(err)
	}
	cmd.Stderr = cmd.Stdout
	// for windows only
	cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true}
	if err = cmd.Start(); err != nil {
		log.Fatal(err)
	}
	buf := make([]byte, 2048)
	for {
		n, err := stdout.Read(buf)
		if n > 0 && err == nil {
			log.Print(string(buf[0:n]))
		}
		if err != nil {
			break
		}
	}
	log.Println("exec " + ffmpegPath + " ok")
}

var upgrader = websocket.Upgrader{
	HandshakeTimeout: 5 * time.Second,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
	Subprotocols: []string{"null"},
}

func realy(c *gin.Context) {
	uri := c.Request.RequestURI
	log.Println(uri)
	key, exist := c.GetQuery("key")
	if !exist || key == "" {
		log.Println("err stream key")
		return
	}
	streamerMapLock.RLock()
	s := streamerMap[key]
	streamerMapLock.RUnlock()
	if s == nil {
		log.Println("err stream key")
		return
	}

	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		return
	}

	playerKey := conn.RemoteAddr().String()
	log.Println("player connected: " + playerKey + " -> streamer " + s.Key)
	s.PlayerMapLock.Lock()
	s.PlayerMap[playerKey] = &player{
		Key:  playerKey,
		Time: time.Now(),
		Conn: conn,
	}
	s.PlayerMapLock.Unlock()

	if s.Lazy {
		go s.run()
	}

	defer func(conn *websocket.Conn) {
		playerKey := conn.RemoteAddr().String()
		log.Println("conn closed: " + playerKey)
		for _, s := range streamerMap {
			if s != nil {
				s.PlayerMapLock.Lock()
				delete(s.PlayerMap, playerKey)
				s.PlayerMapLock.Unlock()
				s.LastDisconnectTime = time.Now()
				if s.Lazy {
					if len(s.PlayerMap) == 0 && s.Alive {
						log.Println("streamer " + s.Key + " has no player, waiting for stop")
					}
				}
			}
		}
		conn.Close()
	}(conn)
	for {
		_, _, err := conn.ReadMessage()
		if err != nil {
			break
		}
	}
}

func streamerAdd(c *gin.Context) {
	err := checkToken(c)
	if err != nil {
		respErr(err.Error(), c)
		return
	}

	key := c.PostForm("key")
	source := c.PostForm("source")
	resolution := c.PostForm("resolution")
	lazyStr := c.PostForm("lazy")
	key = strings.TrimSpace(key)
	source = strings.TrimSpace(source)
	resolution = strings.TrimSpace(resolution)
	lazy, err := strconv.ParseBool(lazyStr)
	if err != nil {
		log.Println(err)
		lazy = true
	}
	if key == "" {
		respErr("key required", c)
		return
	}
	if source == "" {
		respErr("source required", c)
		return
	}
	streamerMapLock.RLock()
	s := streamerMap[key]
	streamerMapLock.RUnlock()
	if s != nil {
		respErr("same key", c)
		return
	}

	var _s streamer
	db.Where(&streamer{
		Key: key,
	}).First(_s)
	if _s.Key != "" {
		respErr("same key", c)
		return
	}
	s = &streamer{
		Key: key,
	}
	s.Source = source
	s.Resolution = resolution
	s.Lazy = lazy
	err = db.Create(&s).Error
	if err != nil {
		respErr(err.Error(), c)
		return
	}
	s.Alive = false
	s.Cmd = nil
	s.PlayerMap = make(map[string]*player)
	s.PlayerMapLock = &sync.RWMutex{}
	if !s.Lazy {
		go s.run()
	}
	streamerMapLock.Lock()
	streamerMap[key] = s
	streamerMapLock.Unlock()
	respOk(nil, c)
}

func streamerList(c *gin.Context) {
	list := make([]*streamerModel, 0)
	for _, s := range streamerMap {
		players := make([]*player, 0)
		for _, p := range s.PlayerMap {
			players = append(players, p)
		}
		sm := &streamerModel{
			Key:        s.Key,
			Source:     s.Source,
			Resolution: s.Resolution,
			Lazy:       s.Lazy,
			Alive:      s.Alive,
			Players:    players,
		}
		list = append(list, sm)
	}
	respOk(list, c)
}

func streamerDelete(c *gin.Context) {
	err := checkToken(c)
	if err != nil {
		respErr(err.Error(), c)
		return
	}

	key := c.PostForm("key")
	if key == "" {
		respErr("key reqpired", c)
		return
	}
	streamerMapLock.RLock()
	s := streamerMap[key]
	streamerMapLock.RUnlock()
	if s != nil {
		streamerMapLock.Lock()
		delete(streamerMap, key)
		streamerMapLock.Unlock()
		s.stop()
	}
	db.Delete(&streamer{
		Key: key,
	})
	respOk(nil, c)
}

func userLogin(c *gin.Context) {
	username := c.PostForm("username")
	password := c.PostForm("password")
	if username != rootUser {
		respErr("err username", c)
		return
	}
	if password != rootPwd {
		respErr("err password", c)
		return
	}
	now := time.Now()
	_t := &token{
		Username: username,
		Role:     0,
		Token:    _md5(username + "_" + now.String()),
		Time:     now,
	}
	t, err := json.Marshal(_t)
	b64 := base64.StdEncoding.EncodeToString([]byte(t))
	if err != nil {
		respErr(err.Error(), c)
	}
	tokenMapLock.Lock()
	tokenMap[username] = _t
	tokenMapLock.Unlock()
	respOk(b64, c)
}

func userLogout(c *gin.Context) {
	t := c.GetHeader("token")
	if t != "" {
		bytes, _ := base64.StdEncoding.DecodeString(t)
		var _t token
		err := json.Unmarshal(bytes, &_t)
		if err == nil {
			tokenMapLock.Lock()
			delete(tokenMap, _t.Username)
			tokenMapLock.Unlock()
		}
	}
	respOk(nil, c)
}

func userResetPwd(c *gin.Context) {
	err := checkToken(c)
	if err != nil {
		respErr(err.Error(), c)
		return
	}
	password := c.PostForm("password")
	if password == "" {
		respErr("password required", c)
		return
	}
	var u user
	db.Model(&u).Update("password", password)
	rootPwd = password
	respOk(nil, c)
}

func getToken(c *gin.Context) *token {
	t := c.GetHeader("_t")
	if t == "" {
		t, _ = c.Cookie("_t")
	}
	if t == "" {
		return nil
	}
	bytes, err := base64.StdEncoding.DecodeString(t)
	if err != nil {
		return nil
	}
	var _t token
	err = json.Unmarshal(bytes, &_t)
	if err != nil {
		return nil
	}
	return &_t
}

func checkToken(c *gin.Context) (err error) {
	_t := getToken(c)
	if _t == nil {
		return errors.New("get token err")
	}
	tokenMapLock.RLock()
	_t1 := tokenMap[_t.Username]
	tokenMapLock.RUnlock()
	if _t1 == nil {
		return errors.New("token expired")
	}
	if _t.Token != _t1.Token {
		return errors.New("token err")
	}
	return nil
}

func respErr(msg string, c *gin.Context) {
	c.JSON(200, gin.H{
		"code": -1,
		"msg":  msg,
	})
}

func respOk(data interface{}, c *gin.Context) {
	c.JSON(200, gin.H{
		"code": 0,
		"msg":  "ok",
		"data": data,
	})
}

type player struct {
	Key  string          `json:"key"`
	Time time.Time       `json:"time"`
	Conn *websocket.Conn `json:"-"`
}

type streamer struct {
	Key                string             `gorm:"column:id;primary_key;NOT NULL;" json:"key"`
	Source             string             `gorm:"type:varchar;NOT NULL;" json:"source"`
	Resolution         string             `gorm:"type:varchar;NOT NULL;" json:"resolution"`
	Lazy               bool               `gorm:"type:tinyint;NOT NULL;" json:"lazy"`
	Alive              bool               `gorm:"-" json:"alive"`
	Cmd                *exec.Cmd          `gorm:"-" json:"-"`
	PlayerMap          map[string]*player `gorm:"-" json:"playerMap"`
	PlayerMapLock      *sync.RWMutex      `gorm:"-" json:"-"`
	LastDisconnectTime time.Time          `gorm:"-" json:"lastDisconnectTime"`
}

func (s *streamer) run() {
	if s.Alive {
		return
	}
	s.Alive = true

	params := []string{ffmpegPath}
	isFile := pathExist(s.Source)
	if !isFile {
		params = append(params, []string{"-rtsp_transport", "tcp", "-stimeout", "5000000"}...)
	}
	params = append(params, []string{"-re", "-i", s.Source,
		"-f", "mpegts", "-codec:v", "mpeg1video", "-preset", "fast", "-nostats", "-r", "24", "-b:v", "700k"}...)
	if s.Resolution != "" {
		params = append(params, []string{"-s", s.Resolution}...)
	}
	params = append(params, []string{"-", "-loglevel", "error"}...)

	cmd := exec.Command(params[0], params[1:]...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Println(err)
		s.Alive = false
		return
	}
	defer stdout.Close()

	// for windows only
	cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true}
	if err = cmd.Start(); err != nil {
		log.Println(err)
		s.Alive = false
		return
	}

	s.Cmd = cmd

	buf := make([]byte, 1024)
	for {
		n, err := stdout.Read(buf)
		if n > 0 && err == nil {
			for _, p := range s.PlayerMap {
				if p != nil && p.Conn != nil {
					p.Conn.WriteMessage(websocket.BinaryMessage, buf[0:n])
				}
			}
		}
		if err != nil || n <= 0 {
			log.Println(err)
			break
		}
	}

	if err = s.Cmd.Wait(); err != nil {
		s.Alive = false
		log.Println("streamer cmd kill err: " + err.Error())
	}
}

func (s *streamer) stop() {
	if s != nil {
		// if s.Cmd != nil && s.Alive {
		if s.Cmd != nil && s.Cmd.Process != nil {
			err := s.Cmd.Process.Kill()
			if err != nil {
				log.Println("streamer cmd kill err: " + err.Error())
			}
		}
		s.Alive = false
	}
}

type user struct {
	Username string `gorm:"column:username;primary_key;NOT NULL;" json:"username"`
	Password string `gorm:"column:password;NOT NULL;" json:"password"`
}

type token struct {
	Username string    `json:"username"`
	Role     int       `json:"role"`
	Token    string    `json:"token"`
	Time     time.Time `json:"time"`
}

type streamerModel struct {
	Key        string    `json:"key"`
	Source     string    `json:"source"`
	Resolution string    `json:"resolution"`
	Lazy       bool      `json:"lazy"`
	Alive      bool      `json:"alive"`
	Players    []*player `json:"players"`
}

func pathExist(path string) bool {
	_, err := os.Stat(path)
	if err == nil {
		return true
	}
	if os.IsNotExist(err) {
		return false
	}
	return false
}

func _md5(str string) string {
	h := md5.New()
	h.Write([]byte(str))
	return hex.EncodeToString(h.Sum(nil))
}
